from torchvision.models.optical_flow.raft import BottleneckBlock
from torchvision.models.resnet import BasicBlock

from layer.PolygonEncoder import PolygonEncoder, VeerCNNPolygonEncoder, NUFTIFFTPolygonEncoder, \
    NUFTIFFTPCAMLPPolygonEncoder, ConcatPolygonEncoder, NUFTSpecPoolPolygonEncoder, NUFTPCAPolygonEncoder, \
    NUFTIFFTMLPPolygonEncoder, NUFTPolygonEncoder
from layer.resnet import ResNet1D
from module.module import MultiLayerFeedForwardNN, MultiLayerFeedForwardNNFlexible


def get_nuftifft_mlp_polygon_encoder(args, pgon_enc_type, pgon_embed_dim, extent,
                                     fft_sigma=2.0, eps=1e-6, elem_batch=300):
    pgon_nuft_embed_dim = compute_pgon_nuft_embed_dim(args, pgon_enc_type, args.nuft_freqXY)
    ffn = get_ffn(args,
                  input_dim=pgon_nuft_embed_dim,
                  output_dim=pgon_embed_dim,
                  f_act=args.spa_f_act,
                  context_str="NUFTIFFTMLPPolygonEncoder")

    pgon_enc = NUFTIFFTMLPPolygonEncoder(
        pgon_embed_dim=pgon_embed_dim,
        ffn=ffn,
        extent=extent,
        eps=eps,
        freqXY=args.nuft_freqXY,
        min_freqXY=args.nuft_min_freqXY,
        max_freqXY=args.nuft_max_freqXY,
        mid_freqXY=args.nuft_mid_freqXY,
        freq_init=args.nuft_freq_init,
        j=args.j,
        embed_norm=args.pgon_nuft_embed_norm_type,
        smoothing='gaussian',
        fft_sigma=fft_sigma,
        elem_batch=elem_batch,
        mode='density',
        device=args.device)
    return pgon_enc


def get_nuftifft_pca_polygon_encoder(args, pgon_enc_type, pgon_embed_dim, extent,
                                     fft_sigma=2.0, eps=1e-6, elem_batch=300):
    pgon_nuft_embed_dim = compute_pgon_nuft_embed_dim(args, pgon_enc_type, args.nuft_freqXY)
    ffn = get_ffn(args,
                  input_dim=pgon_nuft_embed_dim,
                  output_dim=pgon_embed_dim,
                  f_act=args.spa_f_act,
                  context_str="NUFTIFFTPCAMLPPolygonEncoder")

    pgon_enc = NUFTIFFTPCAMLPPolygonEncoder(
        pgon_embed_dim=pgon_embed_dim,
        ffn=ffn,
        extent=extent,
        eps=eps,
        freqXY=args.nuft_freqXY,
        min_freqXY=args.nuft_min_freqXY,
        max_freqXY=args.nuft_max_freqXY,
        mid_freqXY=args.nuft_mid_freqXY,
        freq_init=args.nuft_freq_init,
        j=args.j,
        embed_norm=args.pgon_nuft_embed_norm_type,
        smoothing='gaussian',
        fft_sigma=fft_sigma,
        elem_batch=elem_batch,
        mode='density',
        pca_dim=args.nuft_pca_dim,
        pca_mat=args.pca_mat,
        device=args.device)
    return pgon_enc


def get_polygon_encoder(args, pgon_enc_type, spa_enc, pgon_embed_dim,
                        extent=(-1, 1, -1, 1), padding_mode="circle",
                        dropout=0.5,
                        elem_batch=300, fft_sigma=2.0, eps=1e-6, device="cpu"):
    if pgon_enc_type == "resnet":
        pgon_enc = get_resnet_polygon_encoder(args, spa_enc, pgon_embed_dim)

    elif pgon_enc_type == "veercnn":
        pgon_enc = VeerCNNPolygonEncoder(pgon_embed_dim=pgon_embed_dim,
                                         dropout_rate=dropout,
                                         padding_mode=padding_mode,
                                         device=device)

    elif pgon_enc_type == "nuft_ddsl":
        pgon_enc = get_nuft_ddsl_polygon_encoder(args,
                                                 pgon_enc_type,
                                                 pgon_embed_dim,
                                                 extent,
                                                 fft_sigma=fft_sigma, eps=eps,
                                                 elem_batch=elem_batch)

    elif pgon_enc_type == "nuftifft_ddsl":
        pgon_enc = NUFTIFFTPolygonEncoder(
            extent=extent,
            eps=eps,
            freqXY=args.nuft_freqXY,
            j=args.j,
            smoothing=1,
            fft_sigma=fft_sigma,
            elem_batch=300,
            mode='density',
            device=device)

    elif pgon_enc_type == "nuftifft_mlp":
        pgon_enc = get_nuftifft_mlp_polygon_encoder(args,
                                                    pgon_enc_type,
                                                    pgon_embed_dim,
                                                    extent,
                                                    fft_sigma=fft_sigma, eps=eps,
                                                    elem_batch=elem_batch)

    elif pgon_enc_type in "nuftifft_pca":
        pgon_enc = get_nuftifft_pca_polygon_encoder(args,
                                                    pgon_enc_type,
                                                    pgon_embed_dim,
                                                    extent,
                                                    fft_sigma=fft_sigma, eps=eps,
                                                    elem_batch=elem_batch)


    elif pgon_enc_type == "resnet__nuft_ddsl":
        pgon_enc_1 = get_resnet_polygon_encoder(args, spa_enc, pgon_embed_dim // 2)
        pgon_enc_2 = get_nuft_ddsl_polygon_encoder(args,
                                                   pgon_enc_type="nuft_ddsl",
                                                   pgon_embed_dim=pgon_embed_dim - pgon_embed_dim // 2,
                                                   extent=extent,
                                                   fft_sigma=fft_sigma, eps=eps,
                                                   elem_batch=elem_batch)
        pgon_enc = ConcatPolygonEncoder(pgon_embed_dim,
                                        pgon_enc_1=pgon_enc_1,
                                        pgon_enc_2=pgon_enc_2,
                                        device=device)

    elif pgon_enc_type == "nuft_specpool":
        pgon_enc = get_nuft_specpool_polygon_encoder(args,
                                                     pgon_enc_type,
                                                     pgon_embed_dim,
                                                     extent,
                                                     fft_sigma=fft_sigma, eps=eps,
                                                     elem_batch=elem_batch)

    elif pgon_enc_type in "nuft_pca":
        pgon_enc = get_nuft_pca_polygon_encoder(args,
                                                pgon_enc_type,
                                                pgon_embed_dim,
                                                extent,
                                                fft_sigma=fft_sigma, eps=eps,
                                                elem_batch=elem_batch)

    else:
        raise Exception("Unknow PolygonEncoder type")
    return pgon_enc


def get_nuft_pca_polygon_encoder(args, pgon_enc_type, pgon_embed_dim, extent,
                                 fft_sigma=2.0, eps=1e-6, elem_batch=300):
    pgon_nuft_embed_dim = compute_pgon_nuft_embed_dim(args, pgon_enc_type, args.nuft_freqXY)
    ffn = get_ffn(args,
                  input_dim=pgon_nuft_embed_dim,
                  output_dim=pgon_embed_dim,
                  f_act=args.spa_f_act,
                  context_str="NUFTPCAPolygonEncoder")

    pgon_enc = NUFTPCAPolygonEncoder(
        pgon_embed_dim=pgon_embed_dim,
        dropout_rate=args.dropout,
        ffn=ffn,
        extent=extent,
        eps=eps,
        freqXY=args.nuft_freqXY,
        min_freqXY=args.nuft_min_freqXY,
        max_freqXY=args.nuft_max_freqXY,
        mid_freqXY=args.nuft_mid_freqXY,
        freq_init=args.nuft_freq_init,
        j=args.j,
        embed_norm=args.pgon_nuft_embed_norm_type,
        smoothing='gaussian',
        fft_sigma=fft_sigma,
        elem_batch=elem_batch,
        mode='density',
        nuft_pca_dim=args.nuft_pca_dim,
        pca_mat=args.pca_mat,
        device=args.device)
    return pgon_enc


def get_nuft_specpool_polygon_encoder(args, pgon_enc_type, pgon_embed_dim, extent,
                                      fft_sigma=2.0, eps=1e-6, elem_batch=300):
    pgon_nuft_embed_dim = compute_pgon_nuft_embed_dim(args, pgon_enc_type, args.nuft_freqXY)
    ffn = get_ffn(args,
                  input_dim=pgon_nuft_embed_dim,
                  output_dim=pgon_embed_dim,
                  f_act=args.spa_f_act,
                  context_str="NUFTSpecPoolPolygonEncoder")

    pgon_enc = NUFTSpecPoolPolygonEncoder(
        pgon_embed_dim=pgon_embed_dim,
        dropout_rate=args.dropout,
        ffn=ffn,
        extent=extent,
        eps=eps,
        freqXY=args.nuft_freqXY,
        min_freqXY=args.nuft_min_freqXY,
        max_freqXY=args.nuft_max_freqXY,
        mid_freqXY=args.nuft_mid_freqXY,
        freq_init=args.nuft_freq_init,
        j=args.j,
        embed_norm=args.pgon_nuft_embed_norm_type,
        smoothing='gaussian',
        fft_sigma=fft_sigma,
        elem_batch=elem_batch,
        mode='density',
        spec_pool_max_freqXY=args.spec_pool_max_freqXY,
        spec_pool_min_freqXY_ratio=args.spec_pool_min_freqXY_ratio,
        device=args.device)
    return pgon_enc


def get_resnet_block_by_type(resnet_block_type):
    if resnet_block_type == "basic":
        return BasicBlock
    elif resnet_block_type == "bottleneck":
        return BottleneckBlock
    else:
        raise Exception("Unknown ResNet Block type")


def get_resnet_polygon_encoder(args, spa_enc, pgon_embed_dim):
    resnet_block = get_resnet_block_by_type(args.resnet_block_type)
    # assert len(resnet_layers_per_block) == 3

    pgon_seq_enc = ResNet1D(block=resnet_block,
                            num_layer_list=args.resnet_layers_per_block,
                            in_channels=args.spa_embed_dim,
                            out_channels=pgon_embed_dim,
                            add_middle_pool=args.resnet_add_middle_pool,
                            final_pool=args.resnet_fl_pool_type,
                            padding_mode=args.padding_mode,
                            dropout_rate=args.dropout)

    pgon_enc = PolygonEncoder(spa_enc=spa_enc,
                              pgon_seq_enc=pgon_seq_enc,
                              spa_embed_dim=args.spa_embed_dim,
                              pgon_embed_dim=pgon_embed_dim,
                              device=args.device)
    return pgon_enc


def get_nuft_ddsl_polygon_encoder(args, pgon_enc_type, pgon_embed_dim, extent,
                                  fft_sigma=2.0, eps=1e-6, elem_batch=300):
    pgon_nuft_embed_dim = compute_pgon_nuft_embed_dim(args, pgon_enc_type, args.nuft_freqXY)
    ffn = get_ffn(args,
                  input_dim=pgon_nuft_embed_dim,
                  output_dim=pgon_embed_dim,
                  f_act=args.spa_f_act,
                  context_str="NUFTPolygonEncoder")

    pgon_enc = NUFTPolygonEncoder(
        pgon_embed_dim=pgon_embed_dim,
        dropout_rate=args.dropout,
        ffn=ffn,
        extent=extent,
        eps=eps,
        freqXY=args.nuft_freqXY,
        min_freqXY=args.nuft_min_freqXY,
        max_freqXY=args.nuft_max_freqXY,
        mid_freqXY=args.nuft_mid_freqXY,
        freq_init=args.nuft_freq_init,
        j=args.j,
        embed_norm=args.pgon_nuft_embed_norm_type,
        smoothing='gaussian',
        fft_sigma=fft_sigma,
        elem_batch=elem_batch,
        mode='density',
        device=args.device)
    return pgon_enc


def compute_pgon_nuft_embed_dim(args, pgon_enc_type, freqXY):
    assert len(freqXY) == 2
    fx, fy = freqXY
    if pgon_enc_type == "nuft_ddsl":
        pgon_nuft_embed_dim = fx * (fy // 2 + 1) * 1 * 2
    elif pgon_enc_type == "nuft_specpool":
        assert len(args.spec_pool_max_freqXY) == 2
        wx, wy = args.spec_pool_max_freqXY
        # pgon_nuft_embed_dim = fx * math.ceil(fy/2) * 1 * 2
        pgon_nuft_embed_dim = 4 * wx * wy
    elif pgon_enc_type in ["nuft_pca", "nuftifft_pca"]:
        pgon_nuft_embed_dim = args.nuft_pca_dim
    elif pgon_enc_type == "nuftifft_mlp":
        pgon_nuft_embed_dim = fx * fy
    return pgon_nuft_embed_dim


def get_ffn(args, input_dim, output_dim, f_act, context_str=""):
    if "ffn_type" in args and "ffn_hidden_layers" in args:
        if args.ffn_type == "ffn":
            return MultiLayerFeedForwardNN(
                input_dim=input_dim,
                output_dim=output_dim,
                num_hidden_layers=args.num_hidden_layer,
                dropout_rate=args.dropout,
                hidden_dim=args.hidden_dim,
                activation=f_act,
                use_layernormalize=args.use_layn,
                skip_connection=args.skip_connection,
                context_str=context_str)
        elif args.ffn_type == "ffnf":
            return MultiLayerFeedForwardNNFlexible(
                input_dim=input_dim,
                output_dim=output_dim,
                hidden_layers=args.ffn_hidden_layers,
                dropout_rate=args.dropout,
                activation=f_act,
                use_layernormalize=args.use_layn,
                skip_connection=args.skip_connection,
                context_str=context_str)
        else:
            raise Exception(f"ffn_type: {args.ffn_type} is not defined")
    else:
        return MultiLayerFeedForwardNN(
            input_dim=input_dim,
            output_dim=output_dim,
            num_hidden_layers=args.num_hidden_layer,
            dropout_rate=args.dropout,
            hidden_dim=args.hidden_dim,
            activation=f_act,
            use_layernormalize=args.use_layn,
            skip_connection=args.skip_connection,
            context_str=context_str)
